"""
Unified Reve API Node for ComfyUI

This module implements a single unified node that handles all three Reve API operations
(Create, Edit, Remix) with dynamic inputs based on the selected operation type.
"""

import os
import sys
import logging
import time
import json
from typing import Dict, Any, List, Tuple, Optional

# Add the parent directory to the path to import our modules
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.insert(0, parent_dir)

try:
    from core.reve_api_client import ReveAPIClient, ReveResponse, ReveError
    from core.image_utils import (
        prepare_image_for_api, 
        prepare_multiple_images_for_api,
        base64_to_tensor,
        ImageConversionError
    )
except ImportError as e:
    print(f"Failed to import Reve API modules: {e}")
    raise

import torch

# Configure logging
logger = logging.getLogger(__name__)

# Configure verbose logging for Reve API operations
verbose_logger = logging.getLogger(f"{__name__}.verbose")
verbose_logger.setLevel(logging.INFO)

# Create formatter for verbose logs
verbose_formatter = logging.Formatter(
    '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# Only add handler if it doesn't exist to avoid duplicates
if not verbose_logger.handlers:
    # Create console handler for verbose logging
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(verbose_formatter)
    verbose_logger.addHandler(console_handler)

    # Prevent propagation to avoid duplicate logs
    verbose_logger.propagate = False


class ReveAPI:
    """
    Unified Reve API node supporting Create, Edit, and Remix operations
    with dynamic input system and consistent output handling.
    """
    
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "operation": (["create", "edit", "remix"], {
                    "default": "create",
                    "tooltip": "Select the Reve API operation to perform"
                }),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "A beautiful landscape",
                    "tooltip": "Text description for the operation"
                }),
                "api_key": ("STRING", {
                    "default": "",
                    "tooltip": "Reve API key - paste your API key here"
                })
            },
            "optional": {
                "version": (["latest", "reve-create@20250915", "reve-edit@20250915", "reve-remix@20250915"], {
                    "default": "latest",
                    "tooltip": "Model version to use"
                }),
                "aspect_ratio": (["16:9", "9:16", "3:2", "2:3", "4:3", "3:4", "1:1"], {
                    "default": "3:2",
                    "tooltip": "Output image aspect ratio (for Create and Remix operations)"
                }),
                "output_format": (["png", "jpeg", "webp"], {
                    "default": "png",
                    "tooltip": "Output image format"
                }),
                "reference_image": ("IMAGE", {
                    "tooltip": "Reference image for Edit operation"
                }),
                "reference_images": ("IMAGE", {
                    "tooltip": "Reference images for Remix operation (1-4 images)"
                }),
                "verbose_logging": ("BOOLEAN", {
                    "default": False,
                    "tooltip": "Enable verbose logging for API interactions, performance metrics, and debugging"
                })
            }
        }
    
    RETURN_TYPES = ("IMAGE", "STRING")
    RETURN_NAMES = ("image", "metadata")
    FUNCTION = "process_request"
    CATEGORY = "image/ai"
    DESCRIPTION = "Unified Reve API node for Create, Edit, and Remix operations with dynamic inputs"
    
    def __init__(self):
        self._client = None

    def _log_verbose(self, message: str, data: Dict[str, Any] = None, verbose_enabled: bool = False):
        """Log verbose information if verbose logging is enabled."""
        if verbose_enabled:
            if data:
                formatted_data = json.dumps(data, indent=2, default=str)
                verbose_logger.info(f"{message}\n{formatted_data}")
            else:
                verbose_logger.info(message)

    def _log_api_interaction(self, operation: str, request_params: Dict[str, Any],
                           response: 'ReveResponse', processing_time: float, verbose_enabled: bool = False):
        """Log detailed API interaction information."""
        if not verbose_enabled:
            return

        interaction_data = {
            "operation": operation,
            "request_id": response.request_id,
            "credits_used": response.credits_used,
            "credits_remaining": response.credits_remaining,
            "processing_time_seconds": round(processing_time, 3),
            "success": response.success,
            "request_parameters": {
                "prompt_length": len(request_params.get("prompt", "")),
                "version": request_params.get("version", "unknown"),
                "aspect_ratio": request_params.get("aspect_ratio", "unknown"),
                "has_reference_image": "reference_image" in request_params,
                "reference_images_count": len(request_params.get("reference_images", [])) if "reference_images" in request_params else 0
            }
        }

        if response.error_message:
            interaction_data["error"] = {
                "code": response.error_code,
                "message": response.error_message
            }

        self._log_verbose(f"🔍 Reve API {operation.upper()} Operation Complete", interaction_data, verbose_enabled)
    
    def _get_client(self, api_key: str = "") -> ReveAPIClient:
        """Get or create API client instance."""
        if self._client is None:
            self._client = ReveAPIClient(api_key=api_key if api_key else None)
        return self._client
    
    def _validate_inputs(self, operation: str, **kwargs) -> Dict[str, Any]:
        """Validate inputs based on operation type."""
        errors = []
        
        if operation == "edit":
            if "reference_image" not in kwargs or kwargs["reference_image"] is None:
                errors.append("Edit operation requires a reference_image")
        
        elif operation == "remix":
            if "reference_images" not in kwargs or kwargs["reference_images"] is None:
                errors.append("Remix operation requires reference_images")
            elif hasattr(kwargs["reference_images"], "shape"):
                # Check if it's a batch of images
                batch_size = kwargs["reference_images"].shape[0]
                if batch_size < 1 or batch_size > 4:
                    errors.append(f"Remix operation requires 1-4 reference images, got {batch_size}")
        
        if errors:
            raise ValueError("; ".join(errors))
        
        return kwargs
    
    def _prepare_create_request(self, prompt: str, aspect_ratio: str, version: str, 
                               output_format: str) -> Dict[str, Any]:
        """Prepare request parameters for Create operation."""
        return {
            "prompt": prompt,
            "aspect_ratio": aspect_ratio,
            "version": version,
            "output_format": output_format
        }
    
    def _prepare_edit_request(self, prompt: str, reference_image: torch.Tensor,
                             version: str, output_format: str) -> Dict[str, Any]:
        """Prepare request parameters for Edit operation."""
        try:
            # Convert reference image to base64
            base64_image = prepare_image_for_api(reference_image, output_format)

            return {
                "prompt": prompt,  # This will be converted to 'edit_instruction' in the API client
                "reference_image": base64_image,
                "version": version
            }
        except ImageConversionError as e:
            raise ValueError(f"Failed to prepare reference image: {str(e)}")
    
    def _prepare_remix_request(self, prompt: str, reference_images: torch.Tensor,
                              aspect_ratio: str, version: str, output_format: str) -> Dict[str, Any]:
        """Prepare request parameters for Remix operation."""
        try:
            # Handle batch of images
            if reference_images.dim() == 4:
                # Split batch into individual images
                image_list = [reference_images[i:i+1] for i in range(reference_images.shape[0])]
            else:
                # Single image
                image_list = [reference_images.unsqueeze(0)]
            
            # Convert all images to base64
            base64_images = prepare_multiple_images_for_api(image_list, output_format)
            
            return {
                "prompt": prompt,
                "reference_images": base64_images,
                "aspect_ratio": aspect_ratio,
                "version": version
            }
        except ImageConversionError as e:
            raise ValueError(f"Failed to prepare reference images: {str(e)}")
    
    def _process_response(self, response: ReveResponse) -> Tuple[torch.Tensor, str]:
        """Process API response and convert to ComfyUI outputs."""
        if not response.success:
            error_msg = f"Reve API Error: {response.error_message}"
            if response.error_code:
                error_msg = f"{response.error_code}: {response.error_message}"
            raise RuntimeError(error_msg)
        
        try:
            # Convert base64 image to tensor
            if response.image:
                image_tensor = base64_to_tensor(response.image)
            else:
                raise ValueError("No image data in response")
            
            # Prepare metadata string (enhanced with additional information)
            metadata = {
                "operation": response.metadata.get("operation", "unknown") if response.metadata else "unknown",
                "model_version": response.metadata.get("model_version", "unknown") if response.metadata else "unknown",
                "processing_time": response.metadata.get("processing_time", 0) if response.metadata else 0,
                "parameters": response.metadata.get("parameters", {}) if response.metadata else {},
                "request_id": response.request_id or "unknown",
                "success": response.success
            }
            metadata_str = json.dumps(metadata, indent=2)

            return (
                image_tensor,
                metadata_str
            )
            
        except Exception as e:
            raise RuntimeError(f"Failed to process response: {str(e)}")
    
    def process_request(self, operation: str, prompt: str, api_key: str,
                       version: str = "latest", aspect_ratio: str = "3:2",
                       output_format: str = "png", reference_image: torch.Tensor = None,
                       reference_images: torch.Tensor = None, verbose_logging: bool = False) -> Tuple[torch.Tensor, str]:
        """
        Process Reve API request based on operation type.

        Args:
            operation: Operation type (create, edit, remix)
            prompt: Text description/instruction
            api_key: Reve API key (required)
            version: Model version
            aspect_ratio: Output aspect ratio
            output_format: Output format
            reference_image: Single reference image for edit
            reference_images: Multiple reference images for remix
            verbose_logging: Enable verbose logging for debugging and monitoring

        Returns:
            Tuple of (image_tensor, metadata_str)

        Note:
            When verbose_logging is enabled, additional information including request_id,
            credits_used, credits_remaining, and performance metrics are logged to the console.
        """
        try:
            # Validate API key is provided
            if not api_key or api_key.strip() == "":
                raise ValueError("API key is required. Please paste your Reve API key in the api_key field.")

            # Log start of operation if verbose logging is enabled
            start_time = time.time()
            self._log_verbose(f"🚀 Starting Reve API {operation.upper()} operation", {
                "prompt_length": len(prompt),
                "operation": operation,
                "version": version,
                "aspect_ratio": aspect_ratio,
                "output_format": output_format,
                "has_reference_image": reference_image is not None,
                "reference_images_count": reference_images.shape[0] if reference_images is not None else 0
            }, verbose_logging)

            # Validate inputs based on operation
            inputs = {
                "reference_image": reference_image,
                "reference_images": reference_images
            }
            self._validate_inputs(operation, **inputs)

            # Get API client
            client = self._get_client(api_key)
            
            # Prepare request based on operation type
            api_start_time = time.time()
            request_params = {}

            if operation == "create":
                request_params = {
                    "prompt": prompt,
                    "aspect_ratio": aspect_ratio,
                    "version": version
                }
                response = client.create(**request_params)

            elif operation == "edit":
                request_params = self._prepare_edit_request(
                    prompt, reference_image, version, output_format
                )
                response = client.edit(**request_params)

            elif operation == "remix":
                request_params = self._prepare_remix_request(
                    prompt, reference_images, aspect_ratio, version, output_format
                )
                response = client.remix(**request_params)

            else:
                raise ValueError(f"Unsupported operation: {operation}")

            # Calculate processing time
            processing_time = time.time() - api_start_time
            total_time = time.time() - start_time

            # Log API interaction details
            self._log_api_interaction(operation, request_params, response, processing_time, verbose_logging)

            # Log completion summary
            self._log_verbose(f"✅ Reve API {operation.upper()} operation completed successfully", {
                "total_processing_time_seconds": round(total_time, 3),
                "api_call_time_seconds": round(processing_time, 3),
                "image_processing_time_seconds": round(total_time - processing_time, 3)
            }, verbose_logging)

            # Process and return response
            return self._process_response(response)
            
        except Exception as e:
            # Calculate error time for verbose logging
            error_time = time.time() - start_time if 'start_time' in locals() else 0

            # Log error details if verbose logging is enabled
            self._log_verbose(f"❌ Reve API {operation.upper()} operation failed", {
                "error_type": type(e).__name__,
                "error_message": str(e),
                "operation": operation,
                "time_to_failure_seconds": round(error_time, 3)
            }, verbose_logging)

            logger.error(f"Reve API request failed: {str(e)}")

            # Provide user-friendly error messages for common issues
            error_str = str(e)
            if "API_NOT_AVAILABLE" in error_str or "ENDPOINT_NOT_FOUND" in error_str:
                raise RuntimeError(
                    "❌ Reve API is not currently available.\n\n"
                    "The Reve API endpoints are not accessible. This could be because:\n"
                    "• The API is still in development\n"
                    "• The API requires special access or activation\n"
                    "• The API endpoints have changed\n\n"
                    "Please visit https://api.reve.com or contact Reve support for the latest API status.\n"
                    "You can use the Reve web interface at https://app.reve.com in the meantime."
                )
            else:
                raise RuntimeError(f"Reve API request failed: {str(e)}")
    
    def __del__(self):
        """Cleanup API client on node destruction."""
        if self._client:
            try:
                self._client.close()
            except:
                pass
